{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyscf\n",
    "from pyscf import gto, scf\n",
    "\n",
    "# 构建分子\n",
    "mol = gto.M(\n",
    "    atom='He 0 0 0',\n",
    "    basis='cc-pvdz',\n",
    "    spin=0,\n",
    "    verbose=4,\n",
    ")\n",
    "\n",
    "# 进行HF计算\n",
    "mf = scf.RHF(mol)\n",
    "mf.kernel()\n",
    "\n",
    "# 计算基态能量\n",
    "print('基态能量:', mf.e_tot)\n",
    "\n",
    "# 计算激发态能量\n",
    "import numpy as np\n",
    "from pyscf import gto, scf, fci\n",
    "\n",
    "nelec = mol.nelectrons\n",
    "norb = mol.nao_nr()\n",
    "na = fci.cistring.num_strings(norb, nelec//2)\n",
    "nb = na\n",
    "\n",
    "mycis = fci.addons.cis(mf.mo_coeff, norb, nelec)\n",
    "eci_a, eci_b = mycis.kernel(mf.mo_energy)\n",
    "e_cis = eci_a + eci_b\n",
    "\n",
    "n_roots = 5 # 计算低5个能级\n",
    "n_max_roots = 12\n",
    "conv, e_roots, ci_arr = fci.addons.kernel(mf, norb, nelec, nroots=n_roots, max_cycle=100, max_nroots=n_max_roots)\n",
    "\n",
    "print('激发态能量 (Hartree):')\n",
    "e_roots = e_roots - mf.e_tot\n",
    "for i, e in enumerate(e_roots):\n",
    "    print(f'激发态 {i+1}: {e}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
